{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 简单的扩散模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 线性噪声调度器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    " \n",
    " \n",
    "class LinearNoiseScheduler:\n",
    " \n",
    "    def __init__(self, num_timesteps, beta_start, beta_end):\n",
    "        self.num_timesteps = num_timesteps\n",
    "        self.beta_start = beta_start\n",
    "        self.beta_end = beta_end\n",
    "        \n",
    "        self.betas = torch.linspace(beta_start, beta_end, num_timesteps)\n",
    "        self.alphas = 1. - self.betas\n",
    "        self.alpha_cum_prod = torch.cumprod(self.alphas, dim=0)\n",
    "        self.sqrt_alpha_cum_prod = torch.sqrt(self.alpha_cum_prod)\n",
    "        self.sqrt_one_minus_alpha_cum_prod = torch.sqrt(1 - self.alpha_cum_prod)\n",
    "\n",
    "\n",
    "    def add_noise(self, original, noise, t):\n",
    "        original_shape = original.shape\n",
    "        batch_size = original_shape[0]\n",
    "        \n",
    "        sqrt_alpha_cum_prod = self.sqrt_alpha_cum_prod.to(original.device)[t].reshape(batch_size)\n",
    "        sqrt_one_minus_alpha_cum_prod = self.sqrt_one_minus_alpha_cum_prod.to(original.device)[t].reshape(batch_size)\n",
    "        \n",
    "        # Reshape till (B,) becomes (B,1,1,1) if image is (B,C,H,W)\n",
    "        for _ in range(len(original_shape) - 1):\n",
    "            sqrt_alpha_cum_prod = sqrt_alpha_cum_prod.unsqueeze(-1)\n",
    "        for _ in range(len(original_shape) - 1):\n",
    "            sqrt_one_minus_alpha_cum_prod = sqrt_one_minus_alpha_cum_prod.unsqueeze(-1)\n",
    "        \n",
    "        # Apply and Return Forward process equation\n",
    "        return (sqrt_alpha_cum_prod.to(original.device) * original\n",
    "                + sqrt_one_minus_alpha_cum_prod.to(original.device) * noise)\n",
    "    \n",
    "\n",
    "    def sample_prev_timestep(self, xt, noise_pred, t):\n",
    "        x0 = ((xt - (self.sqrt_one_minus_alpha_cum_prod.to(xt.device)[t] * noise_pred)) /\n",
    "              torch.sqrt(self.alpha_cum_prod.to(xt.device)[t]))\n",
    "        x0 = torch.clamp(x0, -1., 1.)\n",
    "        \n",
    "        mean = xt - ((self.betas.to(xt.device)[t]) * noise_pred) / (self.sqrt_one_minus_alpha_cum_prod.to(xt.device)[t])\n",
    "        mean = mean / torch.sqrt(self.alphas.to(xt.device)[t])\n",
    "        \n",
    "        if t == 0:\n",
    "            return mean, x0\n",
    "        else:\n",
    "            variance = (1 - self.alpha_cum_prod.to(xt.device)[t - 1]) / (1.0 - self.alpha_cum_prod.to(xt.device)[t])\n",
    "            variance = variance * self.betas.to(xt.device)[t]\n",
    "            sigma = variance ** 0.5\n",
    "            z = torch.randn(xt.shape).to(xt.device)\n",
    "            \n",
    "            return mean + sigma * z, x0\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 时间嵌入"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    " \n",
    " \n",
    "def get_time_embedding(time_steps, temb_dim):\n",
    " \n",
    "    assert temb_dim % 2 == 0, \"time embedding dimension must be divisible by 2\"\n",
    "    \n",
    "    # factor = 10000^(2i/d_model)\n",
    "    factor = 10000 ** ((torch.arange(\n",
    "        start=0, end=temb_dim // 2, dtype=torch.float32, device=time_steps.device) / (temb_dim // 2))\n",
    "    )\n",
    "    \n",
    "    # pos / factor\n",
    "    # timesteps B -> B, 1 -> B, temb_dim\n",
    "    t_emb = time_steps[:, None].repeat(1, temb_dim // 2) / factor\n",
    "    t_emb = torch.cat([torch.sin(t_emb), torch.cos(t_emb)], dim=-1)\n",
    "    return t_emb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 下层DownBlock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DownBlock(nn.Module):\n",
    " \n",
    "    def __init__(self, in_channels, out_channels, t_emb_dim,\n",
    "                 down_sample=True, num_heads=4, num_layers=1):\n",
    "        super().__init__()\n",
    "        self.num_layers = num_layers\n",
    "        self.down_sample = down_sample\n",
    "        self.resnet_conv_first = nn.ModuleList(\n",
    "            [\n",
    "                nn.Sequential(\n",
    "                    nn.GroupNorm(8, in_channels if i == 0 else out_channels),\n",
    "                    nn.SiLU(),\n",
    "                    nn.Conv2d(in_channels if i == 0 else out_channels, out_channels,\n",
    "                              kernel_size=3, stride=1, padding=1),\n",
    "                )\n",
    "                for i in range(num_layers)\n",
    "            ]\n",
    "        )\n",
    "        self.t_emb_layers = nn.ModuleList([\n",
    "            nn.Sequential(\n",
    "                nn.SiLU(),\n",
    "                nn.Linear(t_emb_dim, out_channels)\n",
    "            )\n",
    "            for _ in range(num_layers)\n",
    "        ])\n",
    "        self.resnet_conv_second = nn.ModuleList(\n",
    "            [\n",
    "                nn.Sequential(\n",
    "                    nn.GroupNorm(8, out_channels),\n",
    "                    nn.SiLU(),\n",
    "                    nn.Conv2d(out_channels, out_channels,\n",
    "                              kernel_size=3, stride=1, padding=1),\n",
    "                )\n",
    "                for _ in range(num_layers)\n",
    "            ]\n",
    "        )\n",
    "        self.attention_norms = nn.ModuleList(\n",
    "            [nn.GroupNorm(8, out_channels)\n",
    "             for _ in range(num_layers)]\n",
    "        )\n",
    "        \n",
    "        self.attentions = nn.ModuleList(\n",
    "            [nn.MultiheadAttention(out_channels, num_heads, batch_first=True)\n",
    "             for _ in range(num_layers)]\n",
    "        )\n",
    "        self.residual_input_conv = nn.ModuleList(\n",
    "            [\n",
    "                nn.Conv2d(in_channels if i == 0 else out_channels, out_channels, kernel_size=1)\n",
    "                for i in range(num_layers)\n",
    "            ]\n",
    "        )\n",
    "        self.down_sample_conv = nn.Conv2d(out_channels, out_channels,\n",
    "                                          4, 2, 1) if self.down_sample else nn.Identity()\n",
    " \n",
    " \n",
    "    def forward(self, x, t_emb):\n",
    "        out = x\n",
    "        for i in range(self.num_layers):\n",
    "            \n",
    "            # Resnet block of Unet\n",
    "            resnet_input = out\n",
    "            out = self.resnet_conv_first[i](out)\n",
    "            out = out + self.t_emb_layers[i](t_emb)[:, :, None, None]\n",
    "            out = self.resnet_conv_second[i](out)\n",
    "            out = out + self.residual_input_conv[i](resnet_input)\n",
    "            \n",
    "            # Attention block of Unet\n",
    "            batch_size, channels, h, w = out.shape\n",
    "            in_attn = out.reshape(batch_size, channels, h * w)\n",
    "            in_attn = self.attention_norms[i](in_attn)\n",
    "            in_attn = in_attn.transpose(1, 2)\n",
    "            out_attn, _ = self.attentions[i](in_attn, in_attn, in_attn)\n",
    "            out_attn = out_attn.transpose(1, 2).reshape(batch_size, channels, h, w)\n",
    "            out = out + out_attn\n",
    "            \n",
    "        out = self.down_sample_conv(out)\n",
    "        return out\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 中间MidBlock类块"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MidBlock(nn.Module):\n",
    " \n",
    "    def __init__(self, in_channels, out_channels, t_emb_dim, num_heads=4, num_layers=1):\n",
    "        super().__init__()\n",
    "        self.num_layers = num_layers\n",
    "        self.resnet_conv_first = nn.ModuleList(\n",
    "            [\n",
    "                nn.Sequential(\n",
    "                    nn.GroupNorm(8, in_channels if i == 0 else out_channels),\n",
    "                    nn.SiLU(),\n",
    "                    nn.Conv2d(in_channels if i == 0 else out_channels, out_channels, kernel_size=3, stride=1,\n",
    "                              padding=1),\n",
    "                )\n",
    "                for i in range(num_layers+1)\n",
    "            ]\n",
    "        )\n",
    "        self.t_emb_layers = nn.ModuleList([\n",
    "            nn.Sequential(\n",
    "                nn.SiLU(),\n",
    "                nn.Linear(t_emb_dim, out_channels)\n",
    "            )\n",
    "            for _ in range(num_layers + 1)\n",
    "        ])\n",
    "        self.resnet_conv_second = nn.ModuleList(\n",
    "            [\n",
    "                nn.Sequential(\n",
    "                    nn.GroupNorm(8, out_channels),\n",
    "                    nn.SiLU(),\n",
    "                    nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),\n",
    "                )\n",
    "                for _ in range(num_layers+1)\n",
    "            ]\n",
    "        )\n",
    "        \n",
    "        self.attention_norms = nn.ModuleList(\n",
    "            [nn.GroupNorm(8, out_channels)\n",
    "                for _ in range(num_layers)]\n",
    "        )\n",
    "        \n",
    "        self.attentions = nn.ModuleList(\n",
    "            [nn.MultiheadAttention(out_channels, num_heads, batch_first=True)\n",
    "                for _ in range(num_layers)]\n",
    "        )\n",
    "        self.residual_input_conv = nn.ModuleList(\n",
    "            [\n",
    "                nn.Conv2d(in_channels if i == 0 else out_channels, out_channels, kernel_size=1)\n",
    "                for i in range(num_layers+1)\n",
    "            ]\n",
    "        )\n",
    "    \n",
    "    def forward(self, x, t_emb):\n",
    "        out = x\n",
    "        \n",
    "        # First resnet block\n",
    "        resnet_input = out\n",
    "        out = self.resnet_conv_first[0](out)\n",
    "        out = out + self.t_emb_layers[0](t_emb)[:, :, None, None]\n",
    "        out = self.resnet_conv_second[0](out)\n",
    "        out = out + self.residual_input_conv[0](resnet_input)\n",
    "        \n",
    "        for i in range(self.num_layers):\n",
    "            \n",
    "            # Attention Block\n",
    "            batch_size, channels, h, w = out.shape\n",
    "            in_attn = out.reshape(batch_size, channels, h * w)\n",
    "            in_attn = self.attention_norms[i](in_attn)\n",
    "            in_attn = in_attn.transpose(1, 2)\n",
    "            out_attn, _ = self.attentions[i](in_attn, in_attn, in_attn)\n",
    "            out_attn = out_attn.transpose(1, 2).reshape(batch_size, channels, h, w)\n",
    "            out = out + out_attn\n",
    "            \n",
    "            # Resnet Block\n",
    "            resnet_input = out\n",
    "            out = self.resnet_conv_first[i+1](out)\n",
    "            out = out + self.t_emb_layers[i+1](t_emb)[:, :, None, None]\n",
    "            out = self.resnet_conv_second[i+1](out)\n",
    "            out = out + self.residual_input_conv[i+1](resnet_input)\n",
    "        \n",
    "        return out\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## UpBlock上层类快"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "class UpBlock(nn.Module):\n",
    " \n",
    "    def __init__(self, in_channels, out_channels, t_emb_dim, up_sample=True, num_heads=4, num_layers=1):\n",
    "        super().__init__()\n",
    "        self.num_layers = num_layers\n",
    "        self.up_sample = up_sample\n",
    "        self.resnet_conv_first = nn.ModuleList(\n",
    "            [\n",
    "                nn.Sequential(\n",
    "                    nn.GroupNorm(8, in_channels if i == 0 else out_channels),\n",
    "                    nn.SiLU(),\n",
    "                    nn.Conv2d(in_channels if i == 0 else out_channels, out_channels, kernel_size=3, stride=1,\n",
    "                              padding=1),\n",
    "                )\n",
    "                for i in range(num_layers)\n",
    "            ]\n",
    "        )\n",
    "        self.t_emb_layers = nn.ModuleList([\n",
    "            nn.Sequential(\n",
    "                nn.SiLU(),\n",
    "                nn.Linear(t_emb_dim, out_channels)\n",
    "            )\n",
    "            for _ in range(num_layers)\n",
    "        ])\n",
    "        self.resnet_conv_second = nn.ModuleList(\n",
    "            [\n",
    "                nn.Sequential(\n",
    "                    nn.GroupNorm(8, out_channels),\n",
    "                    nn.SiLU(),\n",
    "                    nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),\n",
    "                )\n",
    "                for _ in range(num_layers)\n",
    "            ]\n",
    "        )\n",
    "        \n",
    "        self.attention_norms = nn.ModuleList(\n",
    "            [\n",
    "                nn.GroupNorm(8, out_channels)\n",
    "                for _ in range(num_layers)\n",
    "            ]\n",
    "        )\n",
    "        \n",
    "        self.attentions = nn.ModuleList(\n",
    "            [\n",
    "                nn.MultiheadAttention(out_channels, num_heads, batch_first=True)\n",
    "                for _ in range(num_layers)\n",
    "            ]\n",
    "        )\n",
    "        self.residual_input_conv = nn.ModuleList(\n",
    "            [\n",
    "                nn.Conv2d(in_channels if i == 0 else out_channels, out_channels, kernel_size=1)\n",
    "                for i in range(num_layers)\n",
    "            ]\n",
    "        )\n",
    "        self.up_sample_conv = nn.ConvTranspose2d(in_channels // 2, in_channels // 2,\n",
    "                                                 4, 2, 1) \\\n",
    "            if self.up_sample else nn.Identity()\n",
    "    \n",
    "    def forward(self, x, out_down, t_emb):\n",
    "        x = self.up_sample_conv(x)\n",
    "        x = torch.cat([x, out_down], dim=1)\n",
    "        \n",
    "        out = x\n",
    "        for i in range(self.num_layers):\n",
    "            resnet_input = out\n",
    "            out = self.resnet_conv_first[i](out)\n",
    "            out = out + self.t_emb_layers[i](t_emb)[:, :, None, None]\n",
    "            out = self.resnet_conv_second[i](out)\n",
    "            out = out + self.residual_input_conv[i](resnet_input)\n",
    "            \n",
    "            batch_size, channels, h, w = out.shape\n",
    "            in_attn = out.reshape(batch_size, channels, h * w)\n",
    "            in_attn = self.attention_norms[i](in_attn)\n",
    "            in_attn = in_attn.transpose(1, 2)\n",
    "            out_attn, _ = self.attentions[i](in_attn, in_attn, in_attn)\n",
    "            out_attn = out_attn.transpose(1, 2).reshape(batch_size, channels, h, w)\n",
    "            out = out + out_attn\n",
    " \n",
    "        return out\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## UNet架构"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Unet(nn.Module):\n",
    " \n",
    "    def __init__(self, model_config):\n",
    "        super().__init__()\n",
    "        im_channels = model_config['im_channels']\n",
    "        self.down_channels = model_config['down_channels']\n",
    "        self.mid_channels = model_config['mid_channels']\n",
    "        self.t_emb_dim = model_config['time_emb_dim']\n",
    "        self.down_sample = model_config['down_sample']\n",
    "        self.num_down_layers = model_config['num_down_layers']\n",
    "        self.num_mid_layers = model_config['num_mid_layers']\n",
    "        self.num_up_layers = model_config['num_up_layers']\n",
    "        \n",
    "        assert self.mid_channels[0] == self.down_channels[-1]\n",
    "        assert self.mid_channels[-1] == self.down_channels[-2]\n",
    "        assert len(self.down_sample) == len(self.down_channels) - 1\n",
    "        \n",
    "        # Initial projection from sinusoidal time embedding\n",
    "        self.t_proj = nn.Sequential(\n",
    "            nn.Linear(self.t_emb_dim, self.t_emb_dim),\n",
    "            nn.SiLU(),\n",
    "            nn.Linear(self.t_emb_dim, self.t_emb_dim)\n",
    "        )\n",
    " \n",
    "        self.up_sample = list(reversed(self.down_sample))\n",
    "        self.conv_in = nn.Conv2d(im_channels, self.down_channels[0], kernel_size=3, padding=(1, 1))\n",
    "        \n",
    "        self.downs = nn.ModuleList([])\n",
    "        for i in range(len(self.down_channels)-1):\n",
    "            self.downs.append(DownBlock(self.down_channels[i], self.down_channels[i+1], self.t_emb_dim,\n",
    "                                        down_sample=self.down_sample[i], num_layers=self.num_down_layers))\n",
    "        \n",
    "        self.mids = nn.ModuleList([])\n",
    "        for i in range(len(self.mid_channels)-1):\n",
    "            self.mids.append(MidBlock(self.mid_channels[i], self.mid_channels[i+1], self.t_emb_dim,\n",
    "                                      num_layers=self.num_mid_layers))\n",
    "        \n",
    "        self.ups = nn.ModuleList([])\n",
    "        for i in reversed(range(len(self.down_channels)-1)):\n",
    "            self.ups.append(UpBlock(self.down_channels[i] * 2, self.down_channels[i-1] if i != 0 else 16,\n",
    "                                    self.t_emb_dim, up_sample=self.down_sample[i], num_layers=self.num_up_layers))\n",
    "        \n",
    "        self.norm_out = nn.GroupNorm(8, 16)\n",
    "        self.conv_out = nn.Conv2d(16, im_channels, kernel_size=3, padding=1)\n",
    "    \n",
    "    def forward(self, x, t):\n",
    "        # Shapes assuming downblocks are [C1, C2, C3, C4]\n",
    "        # Shapes assuming midblocks are [C4, C4, C3]\n",
    "        # Shapes assuming downsamples are [True, True, False]\n",
    "        # B x C x H x W\n",
    "        out = self.conv_in(x)\n",
    "        # B x C1 x H x W\n",
    "        \n",
    "        # t_emb -> B x t_emb_dim\n",
    "        t_emb = get_time_embedding(torch.as_tensor(t).long(), self.t_emb_dim)\n",
    "        t_emb = self.t_proj(t_emb)\n",
    "        \n",
    "        down_outs = []\n",
    "        \n",
    "        for idx, down in enumerate(self.downs):\n",
    "            down_outs.append(out)\n",
    "            out = down(out, t_emb)\n",
    "        # down_outs  [B x C1 x H x W, B x C2 x H/2 x W/2, B x C3 x H/4 x W/4]\n",
    "        # out B x C4 x H/4 x W/4\n",
    "            \n",
    "        for mid in self.mids:\n",
    "            out = mid(out, t_emb)\n",
    "        # out B x C3 x H/4 x W/4\n",
    "        \n",
    "        for up in self.ups:\n",
    "            down_out = down_outs.pop()\n",
    "            out = up(out, down_out, t_emb)\n",
    "            # out [B x C2 x H/4 x W/4, B x C1 x H/2 x W/2, B x 16 x H x W]\n",
    "        out = self.norm_out(out)\n",
    "        out = nn.SiLU()(out)\n",
    "        out = self.conv_out(out)\n",
    "        # out B x C x H x W\n",
    "        return out\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 数据集（MNIST）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import os\n",
    " \n",
    "import torchvision\n",
    "from PIL import Image\n",
    "from tqdm import tqdm\n",
    "from torch.utils.data.dataloader import DataLoader\n",
    "from torch.utils.data.dataset import Dataset\n",
    " \n",
    " \n",
    "class MnistDataset(Dataset):\n",
    "    def __init__(self, action, im_path='', split=4, im_ext=''):\n",
    "        Dataset.__init__(self)\n",
    "        self.split = split\n",
    "        self.im_ext = im_ext\n",
    "        self.images_labels = self.load_images(im_path)\n",
    "    \n",
    "    def load_images(self, im_path):\n",
    "        # assert os.path.exists(im_path), \"images path {} does not exist\".format(im_path)\n",
    "        # ims = []\n",
    "        # labels = []\n",
    "        # for d_name in tqdm(os.listdir(im_path)):\n",
    "        #     for fname in glob.glob(os.path.join(im_path, d_name, '*.{}'.format(self.im_ext))):\n",
    "        #         ims.append(fname)\n",
    "        #         labels.append(int(d_name))\n",
    "        # print('Found {} images for split {}'.format(len(ims), self.split))\n",
    "        ds = torchvision.datasets.MNIST(root='mnist_datasets', train=True, download=True)\n",
    "        return ds\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.images_labels)\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        im = self.images_labels[index][0]\n",
    "        im_tensor = torchvision.transforms.ToTensor()(im)\n",
    "        \n",
    "        # Convert input to -1 to 1 range.\n",
    "        im_tensor = (2 * im_tensor) - 1\n",
    "        return im_tensor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\n",
    "    'dataset_params': {\n",
    "        'im_path': 'data/train/images'\n",
    "    },\n",
    " \n",
    "    'diffusion_params': {\n",
    "        'num_timesteps' : 1000,\n",
    "        'beta_start' : 0.0001,\n",
    "        'beta_end' : 0.02\n",
    "    },\n",
    " \n",
    "    'model_params': {\n",
    "        'im_channels' : 1,\n",
    "        'im_size' : 28,\n",
    "        'down_channels' : [32, 64, 128, 256],\n",
    "        'mid_channels' : [256, 256, 128],\n",
    "        'down_sample' : [True, True, False],\n",
    "        'time_emb_dim' : 128,\n",
    "        'num_down_layers' : 2,\n",
    "        'num_mid_layers' : 2,\n",
    "        'num_up_layers' : 2,\n",
    "        'num_heads' : 4\n",
    "    },\n",
    "    'train_params': {\n",
    "        'task_name': 'default',\n",
    "        'batch_size': 64,\n",
    "        'num_epochs': 50,\n",
    "        'num_samples' : 100,\n",
    "        'num_grid_rows' : 10,\n",
    "        'lr': 0.0001,\n",
    "        'ckpt_name': 'ddpm_ckpt.pth'\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done Training ...\n",
      "Loading checkpoint as found one\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\PandaKun\\AppData\\Local\\Temp\\ipykernel_32972\\3103428028.py:39: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  model.load_state_dict(torch.load(os.path.join(train_config['task_name'],\n",
      "100%|██████████| 938/938 [04:20<00:00,  3.60it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished epoch:1 | Loss : 0.0306\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [04:20<00:00,  3.60it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished epoch:2 | Loss : 0.0280\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [04:20<00:00,  3.60it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished epoch:3 | Loss : 0.0270\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [04:20<00:00,  3.60it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished epoch:4 | Loss : 0.0257\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [04:20<00:00,  3.60it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished epoch:5 | Loss : 0.0253\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [04:20<00:00,  3.60it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished epoch:6 | Loss : 0.0248\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [04:20<00:00,  3.60it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished epoch:7 | Loss : 0.0245\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  5%|▌         | 51/938 [00:14<04:09,  3.56it/s]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[67], line 75\u001b[0m\n\u001b[0;32m     71\u001b[0m         torch\u001b[38;5;241m.\u001b[39msave(model\u001b[38;5;241m.\u001b[39mstate_dict(), os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mjoin(train_config[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtask_name\u001b[39m\u001b[38;5;124m'\u001b[39m],\n\u001b[0;32m     72\u001b[0m                                                     train_config[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mckpt_name\u001b[39m\u001b[38;5;124m'\u001b[39m]))\n\u001b[0;32m     74\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mDone Training ...\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m---> 75\u001b[0m \u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[1;32mIn[67], line 65\u001b[0m, in \u001b[0;36mtrain\u001b[1;34m(config)\u001b[0m\n\u001b[0;32m     63\u001b[0m     loss \u001b[38;5;241m=\u001b[39m criterion(noise_pred, noise)\n\u001b[0;32m     64\u001b[0m     losses\u001b[38;5;241m.\u001b[39mappend(loss\u001b[38;5;241m.\u001b[39mitem())\n\u001b[1;32m---> 65\u001b[0m     \u001b[43mloss\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     66\u001b[0m     optimizer\u001b[38;5;241m.\u001b[39mstep()\n\u001b[0;32m     67\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mFinished epoch:\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m | Loss : \u001b[39m\u001b[38;5;132;01m{:.4f}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(\n\u001b[0;32m     68\u001b[0m     epoch_idx \u001b[38;5;241m+\u001b[39m \u001b[38;5;241m1\u001b[39m,\n\u001b[0;32m     69\u001b[0m     np\u001b[38;5;241m.\u001b[39mmean(losses),\n\u001b[0;32m     70\u001b[0m ))\n",
      "File \u001b[1;32m~\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python311\\site-packages\\torch\\_tensor.py:581\u001b[0m, in \u001b[0;36mTensor.backward\u001b[1;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[0;32m    571\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_unary(\u001b[38;5;28mself\u001b[39m):\n\u001b[0;32m    572\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(\n\u001b[0;32m    573\u001b[0m         Tensor\u001b[38;5;241m.\u001b[39mbackward,\n\u001b[0;32m    574\u001b[0m         (\u001b[38;5;28mself\u001b[39m,),\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m    579\u001b[0m         inputs\u001b[38;5;241m=\u001b[39minputs,\n\u001b[0;32m    580\u001b[0m     )\n\u001b[1;32m--> 581\u001b[0m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mautograd\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m    582\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgradient\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minputs\u001b[49m\n\u001b[0;32m    583\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32m~\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python311\\site-packages\\torch\\autograd\\__init__.py:347\u001b[0m, in \u001b[0;36mbackward\u001b[1;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[0;32m    342\u001b[0m     retain_graph \u001b[38;5;241m=\u001b[39m create_graph\n\u001b[0;32m    344\u001b[0m \u001b[38;5;66;03m# The reason we repeat the same comment below is that\u001b[39;00m\n\u001b[0;32m    345\u001b[0m \u001b[38;5;66;03m# some Python versions print out the first line of a multi-line function\u001b[39;00m\n\u001b[0;32m    346\u001b[0m \u001b[38;5;66;03m# calls in the traceback and some print out the last line\u001b[39;00m\n\u001b[1;32m--> 347\u001b[0m \u001b[43m_engine_run_backward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m    348\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtensors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    349\u001b[0m \u001b[43m    \u001b[49m\u001b[43mgrad_tensors_\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    350\u001b[0m \u001b[43m    \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    351\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    352\u001b[0m \u001b[43m    \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    353\u001b[0m \u001b[43m    \u001b[49m\u001b[43mallow_unreachable\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[0;32m    354\u001b[0m \u001b[43m    \u001b[49m\u001b[43maccumulate_grad\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[0;32m    355\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32m~\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python311\\site-packages\\torch\\autograd\\graph.py:825\u001b[0m, in \u001b[0;36m_engine_run_backward\u001b[1;34m(t_outputs, *args, **kwargs)\u001b[0m\n\u001b[0;32m    823\u001b[0m     unregister_hooks \u001b[38;5;241m=\u001b[39m _register_logging_hooks_on_whole_graph(t_outputs)\n\u001b[0;32m    824\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m--> 825\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mVariable\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_execution_engine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_backward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# Calls into the C++ engine to run the backward pass\u001b[39;49;00m\n\u001b[0;32m    826\u001b[0m \u001b[43m        \u001b[49m\u001b[43mt_outputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\n\u001b[0;32m    827\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# Calls into the C++ engine to run the backward pass\u001b[39;00m\n\u001b[0;32m    828\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[0;32m    829\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m attach_logging_hooks:\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import yaml\n",
    "import argparse\n",
    "import os\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "from torch.optim import Adam\n",
    "from torch.utils.data import DataLoader\n",
    " \n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    " \n",
    " \n",
    "def train(config):\n",
    "    diffusion_config = config['diffusion_params']\n",
    "    dataset_config = config['dataset_params']\n",
    "    model_config = config['model_params']\n",
    "    train_config = config['train_params']\n",
    "    \n",
    "    # Create the noise scheduler\n",
    "    scheduler = LinearNoiseScheduler(num_timesteps=diffusion_config['num_timesteps'],\n",
    "                                     beta_start=diffusion_config['beta_start'],\n",
    "                                     beta_end=diffusion_config['beta_end'])\n",
    "    \n",
    "    # Create the dataset\n",
    "    mnist = MnistDataset('train', im_path=dataset_config['im_path'])\n",
    "    mnist_loader = DataLoader(mnist, batch_size=train_config['batch_size'], shuffle=True, num_workers=0)\n",
    "    \n",
    "    # Instantiate the model\n",
    "    model = Unet(model_config).to(device)\n",
    "    model.train()\n",
    "    \n",
    "    # Create output directories\n",
    "    if not os.path.exists(train_config['task_name']):\n",
    "        os.mkdir(train_config['task_name'])\n",
    "    \n",
    "    # Load checkpoint if found\n",
    "    if os.path.exists(os.path.join(train_config['task_name'],train_config['ckpt_name'])):\n",
    "        print('Loading checkpoint as found one')\n",
    "        model.load_state_dict(torch.load(os.path.join(train_config['task_name'],\n",
    "                                                      train_config['ckpt_name']), map_location=device))\n",
    "    # Specify training parameters\n",
    "    num_epochs = train_config['num_epochs']\n",
    "    optimizer = Adam(model.parameters(), lr=train_config['lr'])\n",
    "    criterion = torch.nn.MSELoss()\n",
    "    \n",
    "    # Run training\n",
    "    for epoch_idx in range(num_epochs):\n",
    "        losses = []\n",
    "        for im in tqdm(mnist_loader):\n",
    "            optimizer.zero_grad()\n",
    "            im = im.float().to(device)\n",
    "            \n",
    "            # Sample random noise\n",
    "            noise = torch.randn_like(im).to(device)\n",
    "            \n",
    "            # Sample timestep\n",
    "            t = torch.randint(0, diffusion_config['num_timesteps'], (im.shape[0],)).to(device)\n",
    "            \n",
    "            # Add noise to images according to timestep\n",
    "            noisy_im = scheduler.add_noise(im, noise, t)\n",
    "            noise_pred = model(noisy_im, t)\n",
    " \n",
    "            loss = criterion(noise_pred, noise)\n",
    "            losses.append(loss.item())\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "        print('Finished epoch:{} | Loss : {:.4f}'.format(\n",
    "            epoch_idx + 1,\n",
    "            np.mean(losses),\n",
    "        ))\n",
    "        torch.save(model.state_dict(), os.path.join(train_config['task_name'],\n",
    "                                                    train_config['ckpt_name']))\n",
    "    \n",
    "print('Done Training ...')\n",
    "train(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 采样"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\PandaKun\\AppData\\Local\\Temp\\ipykernel_32972\\3312432755.py:43: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  model.load_state_dict(torch.load(os.path.join(train_config['task_name'],\n",
      "1000it [02:41,  6.18it/s]\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import yaml\n",
    "import os\n",
    "from torchvision.utils import make_grid\n",
    "from tqdm import tqdm\n",
    " \n",
    " \n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    " \n",
    " \n",
    "def sample(model, scheduler, train_config, model_config, diffusion_config):\n",
    " \n",
    "    xt = torch.randn((train_config['num_samples'],\n",
    "                      model_config['im_channels'],\n",
    "                      model_config['im_size'],\n",
    "                      model_config['im_size'])).to(device)\n",
    "    for i in tqdm(reversed(range(diffusion_config['num_timesteps']))):\n",
    "        # Get prediction of noise\n",
    "        noise_pred = model(xt, torch.as_tensor(i).unsqueeze(0).to(device))\n",
    "        \n",
    "        # Use scheduler to get x0 and xt-1\n",
    "        xt, x0_pred = scheduler.sample_prev_timestep(xt, noise_pred, torch.as_tensor(i).to(device))\n",
    "        \n",
    "        # Save x0\n",
    "        ims = torch.clamp(xt, -1., 1.).detach().cpu()\n",
    "        ims = (ims + 1) / 2\n",
    "        grid = make_grid(ims, nrow=train_config['num_grid_rows'])\n",
    "        img = torchvision.transforms.ToPILImage()(grid)\n",
    "        if not os.path.exists(os.path.join(train_config['task_name'], 'samples')):\n",
    "            os.mkdir(os.path.join(train_config['task_name'], 'samples'))\n",
    "        img.save(os.path.join(train_config['task_name'], 'samples', 'x0_{}.png'.format(i)))\n",
    "        img.close()\n",
    " \n",
    " \n",
    "def infer(config):\n",
    "    diffusion_config = config['diffusion_params']\n",
    "    model_config = config['model_params']\n",
    "    train_config = config['train_params']\n",
    "    \n",
    "    # Load model with checkpoint\n",
    "    model = Unet(model_config).to(device)\n",
    "    model.load_state_dict(torch.load(os.path.join(train_config['task_name'],\n",
    "                                                  train_config['ckpt_name']), map_location=device))\n",
    "    model.eval()\n",
    "    \n",
    "    # Create the noise scheduler\n",
    "    scheduler = LinearNoiseScheduler(num_timesteps=diffusion_config['num_timesteps'],\n",
    "                                     beta_start=diffusion_config['beta_start'],\n",
    "                                     beta_end=diffusion_config['beta_end'])\n",
    "    with torch.no_grad():\n",
    "        sample(model, scheduler, train_config, model_config, diffusion_config)\n",
    "\n",
    "infer(config)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
